perf(imputation): cache untreated-projection factorization per fit (#141)#562
Conversation
ImputationDiD's conservative-variance untreated projection
v = -A_0 (A_0'[W]A_0)^{-1} A_1'w has a target-invariant design (A_0/A_1
and the factorization of A_0'[W]A_0) and a target-specific RHS (A_1'w),
but previously rebuilt the sparse design and re-spsolve'd it for every
estimand target (overall ATT, each event-study horizon, each group) and
again for the bootstrap precompute.
Split the build (cached) from the solve: factorize (A_0'[W]A_0) once per
fit() via scipy.sparse.linalg.factorized and solve only the per-target
RHS (factorize-once / solve-many), via a fit-local cache shared across
the analytical and bootstrap paths. Collapses O(1+H+G) factorizations
(2x with bootstrap) to a single build per fit. Mirrors the TwoStageDiD
GMM-sandwich factorized + RuntimeError->dense-lstsq pattern.
Bit-identical to the prior per-target spsolve for a single dense RHS
(both use the SuperLU simple driver with the same defaults); proven at
atol=0 across FE-only, covariate, survey-weighted, and bootstrap paths.
No methodology, numerical, or public-API change.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Overall Assessment✅ Looks good — no unmitigated P0/P1 findings. Executive Summary
MethodologyNo findings. Severity: P3 informational Code QualityNo findings. Severity: P3 informational PerformanceNo findings. Severity: P3 informational MaintainabilityNo findings. Severity: P3 informational Tech DebtNo findings. Severity: P3 informational SecurityNo findings. Severity: P3 informational Documentation/TestsNo findings. Severity: P3 informational |
Summary
ImputationDiDconservative-variance: cache the untreated-projection factorization perfit(). The exact imputation projectionv = -A_0 (A_0'[W]A_0)^{-1} A_1'w(BJS 2024 Theorem 3) has a target-invariant design (A_0/A_1/factorization) and a target-specific RHS (A_1'w), but previously rebuilt the sparse design and re-spsolved it for every estimand target (overall ATT, each event-study horizon, each group) and again for the bootstrap precompute._compute_v_untreated_with_covariatesinto_build_untreated_projection(cached) +_solve_untreated_v(per-target): factorize(A_0'[W]A_0)once viascipy.sparse.linalg.factorizedand solve only the per-target RHS (factorize-once / solve-many), via a fit-local cache shared across the analytical and bootstrap paths. CollapsesO(1 + #horizons + #groups)factorizations (2× with bootstrap) to a single build per fit. Mirrors the TwoStageDiD GMM-sandwichfactorized+RuntimeError→dense-lstsqpattern.Methodology references
ImputationDiD— Borusyak, Jaravel & Spiess (2024) Theorem 3 conservative variance (the untreated imputation projection / Supplementary Proposition A3).docs/methodology/REGISTRY.md## ImputationDiD.spsolvefor a single dense RHS (both use the SuperLU simple driver with the same defaults). No methodology/numerical/public-API change; the REGISTRY "Sparse variance solver" note is updated to describe the factorize-once/solve-many implementation.Validation
tests/test_methodology_imputation.py— newTestImputationVarianceFactorizationCache(bit-identity of the cached factorization vsspsolveincl. the survey-weighted branch; cache-reuse is a numerical no-op; projection built exactly once per fit for the analytical, survey, and bootstrap paths; fit-idempotency / no cross-fit leak) and refreshed the singular-Ω₀ fallback test for the new build-timeRuntimeErrormechanism.tests/test_imputation.py— the two sparse-fallback monkeypatch tests now patchsparse_factorized.atol=0across FE-only, covariate, survey-weighted, and bootstrap paths (overall + event-study + group SEs). black / ruff / mypy clean (0 new errors).Security / privacy
🤖 Generated with Claude Code